#!/usr/bin/env python3

import pandas as pd
import numpy as np
from qiutil import *
from pandas.tseries.offsets import MonthEnd
import qittm
from dataclasses import dataclass

# 用来访问原始数据集，并稍作处理
@dataclass
class RawData:
    equity_data: pd.DataFrame
    index_constituent: pd.DataFrame
    income_data: pd.DataFrame
    code_name: pd.DataFrame
    mrt: pd.DataFrame
    mv_data: pd.DataFrame
    sw_ind: pd.DataFrame
    ind1k8: pd.DataFrame

def load_raw_data(path=DEFAULT_DATA_PATH):
    equity_data = pd.read_hdf(path + "/equity_data.h5")
    index_constituent = pd.read_hdf(path + "/index_constituent.h5")
    income_data = pd.read_hdf(path + "/income_data.h5")
    code_name = pd.read_hdf(path + "/code_name.h5")
    mrt = pd.read_hdf(path + "/monthly_return_turnover.h5")
    mv_data = pd.read_hdf(path + "/mv_data.h5")
    sw_ind = pd.read_hdf(path + "/sw_ind.h5")
    return RawData(equity_data, index_constituent, income_data,
                   code_name, mrt, mv_data, sw_ind, pd.DataFrame())

def kaizen_raw_data(raw):
    # ignore irelavant columns
    raw.equity_data = raw.equity_data[['STOCK_CODE', 'REPORT_PERIOD', 'TOTAL_EQUITY_PARENT', 'ANNOUCE_DATE']]
    raw.equity_data.rename(columns={'REPORT_PERIOD':'E_RPT_PERIOD'}, inplace=True)
    raw.income_data.rename(columns={'REPORT_PERIOD':'I_RPT_PERIOD'}, inplace=True)

    # create calculated columns
    raw.mrt['PREV_MON']= pd.to_datetime(raw.mrt['TRADE_DATE']) + MonthEnd(-1)
    raw.mrt['PREV_MON'] = raw.mrt['PREV_MON'].dt.strftime("%Y%m%d")
    __prep_income_data(raw)
    __prep_index(raw)
    __prep_industries(raw)
    __prep_mv(raw)
    __prep_equity(raw)
    print(raw.equity_data.columns)
    return raw

def load_raw_from_h5store(store):
    equity_data = store.get('equity_data')
    index_constituent = store.get("index_constituent")
    ind1k8 = store.get("ind1k8")
    income_data = store.get("income_data")
    code_name = store.get("code_name")
    mrt = store.get("mrt")
    mv_data = store.get("mv_data")
    sw_ind = store.get("sw_ind")
    return RawData(equity_data, index_constituent, income_data,
                   code_name, mrt, mv_data, sw_ind, ind1k8)

def __prep_index(raw):
    raw.index_constituent.OUT_DATE.replace([None], "22620411", inplace=True)
    filter_expr = f'INDEX_CODE == {CHOSEN_INDEXES} & OUT_DATE == "22620411"';
    raw.ind1k8 = raw.index_constituent.query(filter_expr)

def __prep_income_data(raw):
    raw.income_data = raw.income_data[['STOCK_CODE', 'I_RPT_PERIOD', 'NET_PROFIT_GM', 'ANN_DT']]
    raw.income_data['YEAR'] = raw.income_data['I_RPT_PERIOD'].str.slice(0, 4)

def __prep_industries(raw):
    raw.sw_ind.REMOVE_DATE.replace([None], "22620411", inplace=True)
    industries = raw.sw_ind.IND_NAME.unique()
    for industry_name in industries:
        raw.sw_ind[industry_name] = (raw.sw_ind['IND_NAME'] == industry_name).astype(int)

def __prep_mv(raw):
    # 添加各个时间段的结束日期
    df = raw.mv_data.sort_values(by=['STOCK_CODE', 'TRADE_DATE'])
    df.reset_index(inplace=True, drop=True)
    df[['STOCK_CODE2', 'TRADE_DATE2']]  = df.shift(-1)[['STOCK_CODE', 'TRADE_DATE']]

    bool_ind = df['STOCK_CODE'] != df['STOCK_CODE2']

    df.loc[bool_ind, 'TRADE_DATE2'] = "22620411"
    df.drop(columns = ['STOCK_CODE2'], inplace=True)
    raw.mv_data = df

def __prep_equity(raw):
    # 添加各个时间段的结束日期
    raw.equity_data.drop(columns=['BSPAN', 'ANN_SPAN'], errors='ignore', inplace=True)
    df = raw.equity_data.sort_values(by=['STOCK_CODE', 'E_RPT_PERIOD'])
    df.reset_index(drop=True, inplace=True)
    df[['STOCK_CODE2', 'E_RPT_PERIOD2', 'ANN_DATE2']] = df.shift(-1)[['STOCK_CODE', 'E_RPT_PERIOD', 'ANNOUCE_DATE']]

    bool_ind = df['STOCK_CODE'] != df['STOCK_CODE2']
    df.loc[bool_ind, ['E_RPT_PERIOD2', 'ANN_DATE2']] = MAX_END_DATE
    df.drop(columns=["STOCK_CODE2"], inplace=True)
    raw.equity_data = df

def save_raw_to_h5store(raw, store):
    raw.equity_data.to_hdf(store, "equity_data", mode='a')
    raw.mrt.to_hdf(store, "mrt", mode='a')
    raw.income_data.to_hdf(store, "income_data", mode='a')
    raw.index_constituent.to_hdf(store, "index_constituent", mode='a')
    raw.ind1k8.to_hdf(store, "ind1k8", mode='a')
    raw.mv_data.to_hdf(store, "mv_data", mode='a')
    raw.sw_ind.to_hdf(store, 'sw_ind', mode='a')
    raw.code_name.to_hdf(store, 'code_name', mode='a')

# Data Access Facade
def load_raw_from_files(self, path=DEFAULT_DATA_PATH):
    raw = load_raw_data(path)
    return kaizen_raw_data(raw)

def create_ttm_from_raw(raw):
    print(raw.equity_data.columns)
    return qittm.create_ttm_roe(raw)

def load_raw_and_ttm(data_file=DEFAULT_SILVER_FILE):
    store = pd.HDFStore(data_file)
    raw = load_raw_from_h5store(store)
    ttm = qittm.load_ttm_roe_from_h5store(store)
    store.close()
    return raw, ttm

def save_raw_and_ttm(raw, ttm, data_file=DEFAULT_SILVER_FILE):
    store = pd.HDFStore(data_file, mode='w')
    save_raw_to_h5store(raw, store)
    qittm.save_ttm_roe(ttm, store)
    store.close()
